package panda.natalia.gothread.suspend.instrument;


import org.objectweb.asm.*;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.MethodNode;
import panda.natalia.gothread.suspend.GlobalDefine;
import panda.natalia.gothread.suspend.Instrumented;

import java.util.ArrayList;
import java.util.List;

import static panda.natalia.gothread.suspend.GlobalDefine.isYieldMethod;
import static panda.natalia.gothread.suspend.GlobalDefine.toIntArray;
import static panda.natalia.gothread.suspend.instrument.GoThreadMethodInstrumentor.ASMAPI;

/**
 * a simple framework of java fiber
 *
 * @Author: scpanda2007
 * @Date: 2019/12/10 10:48
 */
public class SuspendOffsetsAfterInstrClassVisitor extends ClassVisitor {

    private final MethodDatabase db;
    private String sourceName, className;

    public SuspendOffsetsAfterInstrClassVisitor(ClassVisitor classVisitor, MethodDatabase db) {
        super(ASMAPI, classVisitor);
        this.db = db;
    }

    @Override
    public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
        this.className = name;
        if(version < Opcodes.V1_5) version = Opcodes.V1_5;
        super.visit(version, access, name, signature, superName, interfaces);
    }

    @Override
    public void visitSource(String sourceName, String debug){
        this.sourceName = sourceName;
        super.visitSource(sourceName, debug);
    }

    public MethodVisitor visitMethod(final int acess, final String name, final String desc, final String signature, final String[] exceptions) {
        if((acess & Opcodes.ACC_NATIVE) == 0 && !isYieldMethod(className, name)){
            final MethodNode mn = new MethodNode(acess, name, desc, signature, exceptions);

            final MethodVisitor outMV = super.visitMethod(acess, name, desc, signature, exceptions);

            return new MethodVisitor(ASMAPI, outMV) {

                private Label currLabel = null;
                private int prevOffset = -1;
                private boolean instrumented;
                private boolean optimized = false;
                private int methodStart = -1, methodEnd = -1;
                private int[] suspCallSites = new int[0];
                private String[] suspCallSiteNames = new String[0];

                private List<Integer> suspOffsetAfterInstrL = new ArrayList<>();

                @Override
                public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {

                    if(GlobalDefine.INSTRUMENTED_DESC.equals(descriptor)){
                        instrumented = true;
                        return new AnnotationVisitor(ASMAPI) {
                            @Override
                            public void visit(String name, Object value) {
                                if(null != name){
                                    switch (name){
                                        case Instrumented.FIELD_NAME_METHOD_START:
                                            methodStart = (Integer) value;
                                            break;
                                        case Instrumented.FIELD_NAME_METHOD_END:
                                            methodEnd = (Integer) value;
                                            break;
                                        case Instrumented.FIELD_NAME_METHOD_OPTIMIZED:
                                            optimized = (Boolean) value;
                                            break;
                                        case Instrumented.FIELD_NAME_SUSPENDABLE_CALL_SITES:
                                            suspCallSites = (int[]) value;
                                            break;
                                        case Instrumented.FIELD_NAME_SUSPENDABLE_CALL_SITES_OFFSETS_AFTER_INSTR:
                                            ;
                                            break;
                                        default:
                                            throw new RuntimeException("Unexcepted `@Instrumented` field: " + name);
                                    }
                                }
                            }

                            @Override
                            public AnnotationVisitor visitArray(String name){
                                if(Instrumented.FIELD_NAME_SUSPENDABLE_CALL_SITE_NAMES.equals(name)){
                                    return new AnnotationVisitor(ASMAPI) {
                                        List<String> callSites = new ArrayList<>();

                                        @Override
                                        public void visit(String name, Object value) {
                                            super.visit(name, value);
                                        }

                                        @Override
                                        public void visitEnd() {
                                            suspCallSiteNames = callSites.toArray(new String[0]);
                                        }
                                    };
                                }else{
                                    return super.visitArray(name);
                                }
                            }
                        };
                    }

                    return super.visitAnnotation(descriptor, visible);
                }

                @Override
                public void visitLocalVariable(String name, String desc, String sig, Label lStart, Label lend, int slot){
                    super.visitLocalVariable(name, desc, sig, lStart, lend, slot);
                }

                @Override
                public void visitLabel(Label label){
                    if(instrumented){
                        currLabel = label;
                    }
                    super.visitLabel(label);
                }

                private void addLine(){
                    final int currOffset = (Integer) currLabel.info;
                    if(currOffset > prevOffset){
                        suspOffsetAfterInstrL.add(currOffset);
                        prevOffset = currOffset;
                    }
                }

                @Override
                public void visitMethodInsn(int opcode, String owner, String name, String desc, boolean isInterface){
                    if(instrumented){
                        final int type = AbstractInsnNode.METHOD_INSN;
                        if(InstrumentMethod.isSuspendableCall(db, type, opcode, owner, name, desc) &&
                        !GlobalDefine.STACK_NAME.equals(owner) && currLabel != null && currLabel.info instanceof Integer)
                            addLine();
                    }

                    super.visitMethodInsn(opcode, owner, name, desc, isInterface);
                }

                @Override
                public void visitInvokeDynamicInsn(String name, String desc, Handle handle, Object ...objects){
                    if(instrumented){
                        final int type = AbstractInsnNode.INVOKE_DYNAMIC_INSN;
                        final int opcode = Opcodes.INVOKEDYNAMIC;
                        if(InstrumentMethod.isSuspendableCall(db, type, opcode, handle.getOwner(), name, desc)
                        && !GlobalDefine.STACK_NAME.equals(handle.getOwner()) && currLabel != null && currLabel.info instanceof Integer)
                            addLine();
                    }
                    super.visitInvokeDynamicInsn(name, desc, handle, objects);
                }

                @Override
                public void visitEnd(){
                    if(instrumented){
                        InstrumentMethod.emitInstrumentedAnn(
                                db, outMV, mn, sourceName, className, optimized, methodStart, methodEnd,
                                suspCallSites, suspCallSiteNames, toIntArray(suspOffsetAfterInstrL)
                        );
                    }

                    super.visitEnd();
                }
            };
        }
        return super.visitMethod(acess, name, desc, signature, exceptions);
    }

}
